# -*- coding:utf-8 -*-
from rich import print
from transformers import AutoTokenizer
import numpy as np
import sys
sys.path.append('..')
from pet_config import *

class HardTemplate(object):
    """
    硬模板，人工定义句子和[MASK]之间的位置关系。
    """

    def __init__(self, prompt: str):
        """
        初始化三个基本属性{prompt,input_list,custom_tokens},调用prompt_analysis() 解析模板给三个属性赋值

        """
        self.prompt = prompt
        self.inputs_list = []                       # 根据文字prompt拆分为各part的列表
        self.custom_tokens = set(['MASK'])          # 从prompt中解析出的自定义token集合

        # print("***********")
        # print(f"custom_tokens->{self.custom_tokens}")
        # print(f"inputs_list->{self.inputs_list}")
        # print(f"prompt->{self.prompt}")
        # print("***********")
        self.prompt_analysis()  # 解析prompt模板

    def prompt_analysis(self):
        """
        将prompt文字模板拆解为可映射的数据结构。
        """
        idx = 0
        while idx < len(self.prompt):
            str_part = ''
            if self.prompt[idx] not in ['{', '}']:
                self.inputs_list.append(self.prompt[idx])
            if self.prompt[idx] == '{':                  # 进入自定义字段
                idx += 1
                while self.prompt[idx] != '}':
                    str_part += self.prompt[idx]             # 拼接该自定义字段的值
                    idx += 1
            elif self.prompt[idx] == '}':
                raise ValueError("Unmatched bracket '}', check your prompt.")
            if str_part:
                # print(f"str_part->{str_part}")
                self.inputs_list.append(str_part)
                self.custom_tokens.add(str_part)             # 将所有自定义字段存储，后续会检测输入信息是否完整
                # print(f"custom_tokens->{self.custom_tokens}")
                # print(f"inputs_list->{self.inputs_list}")
                # print(f"prompt->{self.prompt}")
            idx += 1

    def __call__(self,
                 inputs_dict: dict,
                 tokenizer,
                 mask_length,
                 max_seq_len=512):

        outputs = {
            'text': '', 
            'input_ids': [],
            'token_type_ids': [],
            'attention_mask': [],
            'mask_position': []
        }

        str_formated = ''
        for value in self.inputs_list:
            if value in self.custom_tokens:
                if value == 'MASK':
                    str_formated += inputs_dict[value] * mask_length
                else:
                    str_formated += inputs_dict[value]
            else:
                str_formated += value
        encoded = tokenizer(text=str_formated,
                            truncation=True,
                            max_length=max_seq_len,
                            padding='max_length')
        # print(f'encoded--->{encoded}')
        outputs['input_ids'] = encoded['input_ids']
        outputs['token_type_ids'] = encoded['token_type_ids']
        outputs['attention_mask'] = encoded['attention_mask']
        outputs['text'] = ''.join(tokenizer.convert_ids_to_tokens(encoded['input_ids']))
        mask_token_id = tokenizer.convert_tokens_to_ids(['[MASK]'])[0]
        mask_position = np.where(np.array(outputs['input_ids']) == mask_token_id)[0].tolist()
        outputs['mask_position'] = mask_position
        return outputs


